//
//  ahrs.c
//  KFilter
//
//  This file contains the models of UKF
//

#include <stdio.h>
#include "erza_matrix_lib.h"
#include "config_default.h"
#include "ukf.h"
#include "ahrs.h"

#define P_COV  	0.00505f
#define Q_COV 	0.0000505f
#define ACC_COV 0.165f
#define MAG_COV 0.0001f

ukf_predictionModel AHRSPredictionModel;
ukf_observationModel AHRSAccelerometerModel;
ukf_observationModel AHRSMagnetometerModel;
ukf_observationModel AHRSPressureModel;
AHRS_dataModel AHRSDataModel;

/*output, public*/
AHRS_outputModel AHRSOutputModel;

/* calibration parameters */
AHRS_CalibrationParameters AHRSMagCali;
AHRS_CalibrationParameters AHRSAccCali; // TODO

static float32_t exoInput[EXOINPUTDIM];

//Todo: some covariances are sqrt'ed some are not, determine which ones are supposed to be sqrt'ed
int AHRSInit() {
	float32_t covariance[20];

	/*Process model*/
	AHRSPredictionModel.stateDim = 7;
	AHRSPredictionModel.observationCovDim = 6;
	AHRSPredictionModel.exoInputDim = 6;
	AHRSPredictionModel.scaleFactor = 3.0f;

	covariance[0] = P_COV;
	covariance[1] = P_COV;
	covariance[2] = P_COV;
	covariance[3] = P_COV;
	covariance[4] = P_COV;
	covariance[5] = P_COV;
	covariance[6] = P_COV;

	/*Set process noise covariance*/
	setCDKFDiagonalCovariance(&AHRSPredictionModel.processCov, AHRSPredictionModel.stateDim, covariance, AHRSPredictionModel.stateDim, TRUE);

	covariance[0] = Q_COV;
	covariance[1] = Q_COV;
	covariance[2] = Q_COV;
	covariance[3] = Q_COV;
	covariance[4] = Q_COV;
	covariance[5] = Q_COV;

	/*Set observation noise covariance*/
	setCDKFDiagonalCovariance(&AHRSPredictionModel.observationCov, AHRSPredictionModel.observationCovDim, covariance, AHRSPredictionModel.observationCovDim, FALSE);

	AHRSPredictionModel.fx_predictionUpdate = &predictionUpdate;

	/***********************
	 * Measurement Models
	 ***********************/

	/*Measurement Model - Acc*/
	AHRSAccelerometerModel.noiseDim = 3;
	AHRSAccelerometerModel.stateDim = 3;
	AHRSAccelerometerModel.exoInputDim = 0;
	AHRSAccelerometerModel.fx_observationUpdate = &accelerometerObservationUpdate;

	covariance[0] = ACC_COV;
	covariance[1] = ACC_COV;
	covariance[2] = ACC_COV;
	setCDKFDiagonalCovariance(&AHRSAccelerometerModel.cov, AHRSAccelerometerModel.noiseDim, covariance, AHRSAccelerometerModel.noiseDim, FALSE);

	/*Measurement Model - Mag*/
	AHRSMagnetometerModel.noiseDim = 3;
	AHRSMagnetometerModel.stateDim = 3;
	AHRSMagnetometerModel.exoInputDim = 0;
	AHRSMagnetometerModel.fx_observationUpdate = &magnetometerObservationUpdate;
	covariance[0] = MAG_COV;
	covariance[1] = MAG_COV;
	covariance[2] = MAG_COV;
	setCDKFDiagonalCovariance(&AHRSMagnetometerModel.cov, AHRSMagnetometerModel.noiseDim, covariance, AHRSMagnetometerModel.noiseDim, FALSE);

	/*Pressure*/
	AHRSPressureModel.noiseDim = 1;
	AHRSPressureModel.stateDim = 1;
	AHRSPressureModel.exoInputDim = 0;
	AHRSPressureModel.fx_observationUpdate = &barameterObservationUpdate;
	covariance[0] = MAG_COV;
	setCDKFDiagonalCovariance(&AHRSPressureModel.cov, AHRSPressureModel.noiseDim, covariance, AHRSPressureModel.noiseDim, FALSE);

	ukf_Init(&AHRSPredictionModel);

	AHRSPredictionModel.state.pData[0] = 1.0f;

	/****************************************
	 * End UKF prediction and measurment init
	 ****************************************/



	/* magnetometer data init */
	float32_t DEFAULT_IMU_MAG_INCL = -60.98f;
	float32_t DEFAULT_IMU_MAG_DECL = -13.82f;
	float32_t DEG_TO_RAD = PI / 180;

	const float32_t mag1 = cosf(DEFAULT_IMU_MAG_INCL * DEG_TO_RAD);
	const float32_t mag2 = 0.0f;
	const float32_t mag3 = sinf(DEFAULT_IMU_MAG_INCL * DEG_TO_RAD);

	AHRSDataModel.mag0rf = (mag1 * cos(DEFAULT_IMU_MAG_DECL * DEG_TO_RAD) - mag2 * sin(DEFAULT_IMU_MAG_DECL * DEG_TO_RAD));
	AHRSDataModel.mag1rf = -(mag2 * cos(DEFAULT_IMU_MAG_DECL * DEG_TO_RAD) + mag1 * sin(DEFAULT_IMU_MAG_DECL * DEG_TO_RAD));
	AHRSDataModel.mag2rf = mag3;


	/*Init Magnetometer calibration data*/
	AHRSMagCali.sc00 = 0.00200514834530657;
	AHRSMagCali.sc01 = 0.000145660478826823;
	AHRSMagCali.sc02 = -2.28229956453876e-05;
	AHRSMagCali.sc10 = 3.49857099645271e-07;
	AHRSMagCali.sc11 = 0.00219537371085451;
	AHRSMagCali.sc12 = 3.24479067318367e-05;
	AHRSMagCali.sc20 = -5.48178004495311e-08;
	AHRSMagCali.sc21 = 8.13470415420607e-08;
	AHRSMagCali.sc22 = 0.00184810513348282;

	AHRSMagCali.b0 = -42.163146046205;
	AHRSMagCali.b1 = -191.918971854855;
	AHRSMagCali.b2 = -178.53760333671;

	return 1;
}

void processInertialSensors(float32_t gx, float32_t gy, float32_t gz, float32_t ax, float32_t ay, float32_t az) {
	exoInput[GX] = gx;
	exoInput[GY] = gy;
	exoInput[GZ] = gz;
	exoInput[AX] = ax;
	exoInput[AY] = ay;
	exoInput[AZ] = az;

	ukf_processPrediction(&AHRSPredictionModel, exoInput, AHRS_TIMESTAMP);
}

void processAccelerometerObservation(float32_t ax, float32_t ay, float32_t az) {
	float32_t norm = sqrtf(ax*ax + ay*ay + az*az);
	exoInput[OBSx] = ax/norm;
	exoInput[OBSy] = ay/norm;
	exoInput[OBSz] = az/norm;

	ukf_processObservation(&AHRSPredictionModel, &AHRSAccelerometerModel, exoInput);
}

void processBarometricPressureObservation(float32_t p) {
	exoInput[OBSx] = p;

	ukf_processObservation(&AHRSPredictionModel, &AHRSPressureModel, exoInput);
}

void processOutput(uint8_t AHRS_outputType) {
	float32_t normalizedQuat[4];
	float32_t pitch, roll, yaw;



	switch (AHRS_outputType) {
	case AHRS_OUT_QUATERNION:
		/*Set the quaternion in the output*/
		quatNorm(AHRSPredictionModel.state.pData,normalizedQuat);
		quatToEuler(normalizedQuat, &yaw, &pitch, &roll);

		yaw = yaw * 180/PI;
		pitch = pitch*180/PI;
		roll = roll * 180/PI;

		printf("%f %f %f\n", yaw, pitch, roll);

		break;
	case AHRS_OUT_EULER:
		break;
	case AHRS_OUT_ROTATIONMATRIX:
		break;
	}
	//mprintf("final sx", &(AHRSPredictionModel.Sx));

}

/**************************
 * callback functions below...
 **************************/

/*prediction*/
void predictionUpdate(float32_t * inState, float32_t * outState, float32_t * noise, float32_t * exogenousInput, float32_t dt) {
	const float32_t q0 = inState[0];
	const float32_t q1 = inState[1];
	const float32_t q2 = inState[2];
	const float32_t q3 = inState[3];
	const float32_t gyroBiasX = inState[4];
	const float32_t gyroBiasY = inState[5];
	const float32_t gyroBiasZ = inState[6];

	const float32_t gyroNoiseX = noise[0];
	const float32_t gyroNoiseY = noise[1];
	const float32_t gyroNoiseZ = noise[2];

	const float32_t gyroBiasNoiseX = noise[3];
	const float32_t gyroBiasNoiseY = noise[4];
	const float32_t gyroBiasNoiseZ = noise[5];

	// calculate output bias
	outState[4] = gyroBiasX + gyroBiasNoiseX * dt;
	outState[5] = gyroBiasY + gyroBiasNoiseY * dt;
	outState[6] = gyroBiasZ + gyroBiasNoiseZ * dt;

	// rotate input quaternion
	float32_t wx = (exogenousInput[GX] + gyroBiasX + gyroNoiseX) * dt;
	float32_t wy = (exogenousInput[GY] + gyroBiasY + gyroNoiseY) * dt;
	float32_t wz = (exogenousInput[GZ] + gyroBiasZ + gyroNoiseZ) * dt;

	const float32_t s = (sqrtf(wx * wx + wy * wy + wz * wz)) * 0.5f;
	float32_t t = -(0.5f * sinf(s) / s);

	wx *= t;
	wy *= t;
	wz *= t;

	const float32_t qMag = q0 * q0 + q1 * q1 + q2 * q2 + q3 * q3;
	const float32_t lg = cosf(s) + (1 - qMag * qMag) * dt * dt;

	const float32_t q0Out = lg * q0 + wx * q1 + wy * q2 + wz * q3;
	const float32_t q1Out = -wx * q0 + lg * q1 - wz * q2 + wy * q3;
	const float32_t q2Out = -wy * q0 + wz * q1 + lg * q2 - wx * q3;
	const float32_t q3Out = -wz * q0 - wy * q1 + wx * q2 + lg * q3;

	//const float32_t qOutNorm = sqrtf(q0Out * q0Out + q1Out * q1Out + q2Out * q2Out + q3Out * q3Out);

	outState[0] = q0Out; /// qOutNorm;
	outState[1] = q1Out; /// qOutNorm;
	outState[2] = q2Out; /// qOutNorm;
	outState[3] = q3Out; /// qOutNorm;
}

void accelerometerObservationUpdate(float32_t * inState, float32_t * outState, float32_t * noise, float32_t * exogenousInput) {
	const float32_t q0 = inState[0];
	const float32_t q1 = inState[1];
	const float32_t q2 = inState[2];
	const float32_t q3 = inState[3];

	outState[0] = 2.0f * (q1 * q3 - q0 * q2) + noise[0];
	outState[1] = 2.0f * (q2 * q3 + q0 * q1) + noise[1];
	outState[2] = 1.0f - (2.0f * (q1 * q1 + q2 * q2)) + noise[2];
}

void magnetometerObservationUpdate(float32_t * inState, float32_t * outState, float32_t * noise, float32_t * exogenousInput) {

	// reverse quaternion for rotation
	const float32_t w = inState[0];
	const float32_t x = -inState[1];
	const float32_t y = -inState[2];
	const float32_t z = -inState[3];

	const float32_t v0m1 = AHRSDataModel.mag0rf;
	const float32_t v0m2 = AHRSDataModel.mag1rf;
	const float32_t v0m3 = AHRSDataModel.mag2rf;

	// rorate by reverse quaternion
	const float32_t rotateByRevQuat1 = w * w * v0m1 + 2.0 * y * w * v0m3 - 2.0 * z * w * v0m2 + x * x * v0m1 + 2.0 * y * x * v0m2 + 2.0 * z * x * v0m3 - z * z * v0m1 - y * y * v0m1;
	const float32_t rotateByRevQuat2 = 2.0 * x * y * v0m1 + y * y * v0m2 + 2.0 * z * y * v0m3 + 2.0 * w * z * v0m1 - z * z * v0m2 + w * w * v0m2 - 2.0 * x * w * v0m3 - x * x * v0m2;
	const float32_t rotateByRevQuat3 = 2.0 * x * z * v0m1 + 2.0 * y * z * v0m2 + z * z * v0m3 - 2.0 * w * y * v0m1 - y * y * v0m3 + 2.0 * w * x * v0m2 - x * x * v0m3 + w * w * v0m3;

	outState[0] = -rotateByRevQuat2 + noise[0];
	outState[1] = rotateByRevQuat3 + noise[1];
	outState[2] = -rotateByRevQuat1 + noise[2];
}

void barameterObservationUpdate(float32_t * inState, float32_t * outState, float32_t * noise, float32_t * exogenousInput) {
	outState[0] = inState[0] + noise[0];
}
/*Helper functions below*/

void calibrateMag(float *mx, float *my, float *mz){
	float magCaliX = AHRSMagCali.sc00 * (*mx - AHRSMagCali.b0) + AHRSMagCali.sc01 * (*my - AHRSMagCali.b1) + AHRSMagCali.sc02 * (*mz - AHRSMagCali.b2);
	float magCaliX = AHRSMagCali.sc10 * (*mx - AHRSMagCali.b0) + AHRSMagCali.sc11 * (*my - AHRSMagCali.b1) + AHRSMagCali.sc12 * (*mz - AHRSMagCali.b2);
	float magCaliX = AHRSMagCali.sc20 * (*mx - AHRSMagCali.b0) + AHRSMagCali.sc21 * (*my - AHRSMagCali.b1) + AHRSMagCali.sc22 * (*mz - AHRSMagCali.b2);

	float magNorm = sqrtf(magCaliX*magCaliX + magCaliY*magCaliY + magCaliZ*magCaliZ);
	*mx = magCaliX / magNorm;
	*my = magCaliY / magNorm;
	*mz = magCaliZ / magNorm;
}

void calibrateAcc(float *ax, float *ay, float *az){

}

void calibrateGyro(float *gx, float *gy, float *gz){

}
